import os
from collections import deque
from dataclasses import dataclass
from typing import List, Optional

from base import get_base_file_name


class _Node(object):
    """
    词典树的节点
    """
    count = 0  # 节点数

    def __init__(self, key: str, offset: int, level: int) -> None:
        """
        一个节点必要的数据项，如：字符串 ba 的存储方式为：
        第一个节点：b
        逻辑结构：{ key: b, offset: 0, level: 0b_00000000 }
        物理结构: 0x62_00000000_0000，占用的存储空间为：1 + 4 + 2 = 7 个字节
        第二个节点：a
        逻辑结构：{ key: a, offset: 4, level: 0b_00000001 }
        物理结构：0x61_00000004_0001
        :param key: 一个字符
        :param offset: 第一次出现的位置
        :param level: 层级
        """
        super().__init__()
        self.key = key
        self.offset = offset
        self.level = level
        self.children: List[_Node] = []

    def isChildrenEmpty(self):
        return len(self.children) == 0

    @property
    def isLeaf(self):
        return self.offset > 0

    def append(self, node):
        _Node.count += 1
        if self.isChildrenEmpty():
            self.children.append(node)
            return

        if not self.hasChild(node.key):
            self.children.append(node)
            self.children.sort(key=lambda x: x.key)

    @property
    def firstChild(self):
        if len(self.children) > 0:
            return self.children[0]
        return None

    def hasChild(self, key: str):
        return self.findChild(key) is not None

    def findChild(self, key: str):
        if self.isChildrenEmpty():
            return None

        for n in self.children:
            if n.key == key:
                return n
        return None

    def prefixSearch(self, prefix: str, single=True):
        current = self
        path = ''
        for ch in prefix:
            current = current.findChild(ch)
            if current is None:
                break
            else:
                path += ch
            if path == prefix:
                break
        if single:
            return current if path == prefix else None

        rtn = []
        if path == prefix:
            if current.isLeaf:
                rtn.append(current)
            for child in current.children:
                if child.isLeaf:
                    rtn.append(child)
            if len(rtn) == 0:
                parent = current
                child = parent.firstChild
                while child is not None:
                    if child.isLeaf:
                        rtn.append(child)
                        break
                    parent = child
                    child = parent.firstChild
        return rtn

    def findFirstLeaf(self, prefix: str):
        current = self
        path = ''
        for ch in prefix:
            target = current.findChild(ch)
            if target is None:
                break
            else:
                path += ch
                current = target
            if path == prefix:
                break
        if path == prefix and current.isLeaf:
            return current
        leaf = current
        while not current.isChildrenEmpty():
            for n in current.children:
                if n.isLeaf:
                    leaf = n
                    break
            if leaf.isLeaf:
                break
            else:
                current = current.firstChild
        return leaf if leaf.isLeaf else None

    def contains(self, path: str):
        node = self.prefixSearch(path)
        return False if node is None else node.isLeaf

    def dump(self, out):
        # dump self
        if len(self.key) == 0:  # root
            out.write(b'root')
            out.write(_Node.count.to_bytes(4, byteorder=ImDict.kByteOrder))  # 根据 count 值可以算出整棵树占用的空间大小
        else:
            out.write(bytes(self.key, encoding=ImDict.kEncoding))
            out.write(self.offset.to_bytes(4, byteorder=ImDict.kByteOrder))
            out.write(self.level.to_bytes(2, byteorder=ImDict.kByteOrder))
        out.flush()
        # dump children
        for n in self.children:
            n.dump(out)

    def __str__(self) -> str:
        return '<root>' if len(self.key) == 0 else '{} {}@{}'.format(self.key, self.level, self.offset)


class ImDictItem(object):
    """
    词库中的词条
    """

    def __init__(self, text: str, code: str, weight=0) -> None:
        super().__init__()
        self.text = text
        self.code = code
        self.weight = weight

    def __str__(self) -> str:
        if self.weight > 0:
            return '{}\t{}\t{}'.format(self.text, self.code, self.weight)
        return '{}\t{}'.format(self.text, self.code)


class ImDict(object):
    """
    词库文件
    """
    kHead = 'ImDict/UTF-8 Big'
    kVersionName = '1.0'
    kVersionCode = 1_00  # 1.0
    kEncoding = 'utf-8'
    kByteOrder = 'big'

    def __init__(self, src: str, base: int) -> None:
        super().__init__()
        self.src = src
        self.base = base  # 词库内容的基地址
        self.root: Optional[_Node] = None
        self.fp = None  # file pointer
        self.posCache = {}  # 缓存一下文件指针的位置，$prefix: {$start: $pos}
        self.lastLoadCode = None  # 最后一个读到的编码
        self.posCodeMap = {}  # pos -> code

    def useRoot(self, root: _Node):
        self.root = root

    def close(self):
        if self.fp:
            self.fp.close()

    def prefixMatch(self, prefix: str, start=0, limit=9):
        if self.fp is None:
            self.fp = open(self.src, 'rb')
        rtn = []
        if self.root is None:
            return rtn

        pos = -1
        if prefix in self.posCache:
            values = self.posCache.get(prefix)
            if start in values:
                pos = values.get(start)
        if pos == -1:
            self.posCache.clear()
            leaf = self.root.findFirstLeaf(prefix)
            if leaf is not None:
                pos = self.base + leaf.offset
                self.posCache[prefix] = {}
                self.posCache[prefix][0] = pos
        if pos > 0:
            if start == 0:
                self.posCodeMap.clear()
            pos = self._loadItemsTo(pos, prefix, rtn, limit)
            if len(rtn) == limit and pos > 0:
                self.posCache[prefix][start + len(rtn)] = pos
        return rtn

    def perfectMatch(self, code: str):
        if self.fp is None:
            self.fp = open(self.src, 'rb')
        if self.root is None:
            return []

        rtn = []
        node = self.root.prefixSearch(code)
        if node is None:
            return rtn
        rtn.extend(self._loadItems(node))
        return rtn

    def _loadItems(self, node: _Node, limit=1) -> List[ImDictItem]:
        rtn = []
        codeLength = node.level
        pos = self.base + node.offset
        self.fp.seek(pos)
        length = self.fp.read(1)[0]
        if length == codeLength:
            buffer = bytearray()
            code = self.fp.read(length).decode()
            while True:
                first = self.fp.read(1)[0]
                if first == 0:
                    text = buffer.decode(encoding=ImDict.kEncoding)
                    rtn.append(ImDictItem(text, code))
                    buffer.clear()
                    second = self.fp.read(1)[0]
                    if second == 0:
                        break
                    else:
                        buffer.append(second)
                else:
                    buffer.append(first)
        return rtn

    def _loadItemsTo(self, offset: int, prefix: str, result: List[ImDictItem], limit=1) -> int:
        pos = offset
        fp = self.fp
        fp.seek(pos)
        buffer = bytearray()
        count = 0
        # check last code
        if pos in self.posCodeMap:
            code = self.posCodeMap.get(pos)
        else:
            # read code length
            length = fp.read(1)[0]
            print('>>> <{}> offset={}, length={}'.format(__name__, hex(offset), length))
            # read code
            code = '' if length == 0 else fp.read(length).decode()
        if not code.startswith(prefix):
            self.lastLoadCode = None
            return -1
        end = False
        backOff = 0  # 如果多读了一部分字节，需要回退回去
        while count < limit:
            if end:
                length = fp.read(1)[0]
                if length == 0:
                    backOff = -1
                    break
                code = fp.read(length).decode()
                if not code.startswith(prefix):
                    backOff = -(length + 1)
                    break
                end = False
            first = fp.read(1)[0]
            if first == 0:  # 词条分隔标志, 上一个词条结束
                text = buffer.decode(encoding=ImDict.kEncoding)
                result.append(ImDictItem(text, code))
                count += 1
                buffer.clear()
                second = fp.read(1)[0]
                if second == 0:  # 词条结束
                    end = True
                else:  # 正在读取词条
                    buffer.append(second)
            else:  # 正在读取词条
                buffer.append(first)
        endPos = fp.tell() + backOff
        if len(buffer) > 0:  # 存在一个不完整的词条
            endPos -= len(buffer)
            self.posCodeMap[endPos] = code
        return endPos


class ImDictMaker(object):
    """
    词典构建器
    """

    def __init__(self) -> None:
        super().__init__()
        self.src = None
        self.target = None
        self.items = []

    def make(self):
        if len(self.items) == 0:
            raise Exception('Items is empty!')
        self._sort()
        self._compile()

    def readFromSource(self, src):
        self.src = src
        self.target = get_base_file_name(self.src) + '.dic'
        self.items.clear()
        with open(src, 'r', encoding=ImDict.kEncoding) as f:
            for line in f:
                if line.startswith('#'):
                    continue
                segments = line.split('\t', maxsplit=3)
                text = segments[0].strip()
                code = segments[1].strip()
                weight = 0 if len(segments) < 3 else int(segments[2].strip())
                self.items.append(ImDictItem(text, code, weight))
        return self

    def _sort(self):
        self.items.sort(key=lambda x: (x.code, -x.weight))

    def _compile(self):
        with open(self.target, 'wb') as out:
            _writeDictHead(out)
            self._writeContent(out)
            self._merge(out)
        os.remove(self.src + '.tmp')

    def _writeContent(self, out):
        self.root = _Node('', 0, 0)
        root = self.root
        base = out.tell()  # 基地址位置
        out.write(b'base')  # 基地址占位
        self.tmp = open(self.src + '.tmp', 'wb+')  # 以二进制方式打开读写
        tmp = self.tmp
        tmp.write(b'item')
        tmp.flush()
        prevCode = ''
        for item in self.items:
            code = item.code
            if code == prevCode:
                offset = tmp.tell()
            else:
                tmp.write(b'\x00')
                offset = tmp.tell()
                tmp.write(len(code).to_bytes(1, byteorder=ImDict.kByteOrder))
                tmp.write(bytes(code, encoding=ImDict.kEncoding))
                prevCode = code
            tmp.write(bytes(item.text, encoding=ImDict.kEncoding))
            tmp.write(b'\x00')
            tmp.flush()
            current = root.prefixSearch(code)
            if current is not None:  # 已存在
                if not current.isLeaf:
                    current.offset = offset
                continue
            # 不存在
            current = root
            length = len(code)
            for i in range(0, length - 1):
                ch = code[i]
                n = current.findChild(ch)
                if n is None:  # not found
                    child = _Node(ch, 0, i + 1)
                    current.append(child)
                    current = child
                else:
                    current = n
            last = code[-1]
            current.append(_Node(last, offset, length))
        root.dump(out)
        headLength = out.tell()
        out.seek(base)
        out.write(headLength.to_bytes(4, byteorder=ImDict.kByteOrder))
        out.flush()
        out.seek(headLength)
        # 内容结束标志
        tmp.write(b'\x00\x00\x00')
        tmp.flush()

    def _merge(self, out):
        # 合并文件
        tmp = self.tmp
        tmp.seek(0)
        while True:
            content = tmp.read(1024)
            if content:
                out.write(content)
                out.flush()
            else:
                break
        tmp.close()


def _writeDictHead(out):
    out.write(ImDict.kHead.encode())
    out.write(ImDict.kVersionCode.to_bytes(4, byteorder=ImDict.kByteOrder))
    versionName = 'version:' + ImDict.kVersionName
    out.write(versionName.encode())
    # 保证有一个空行
    exists = (4 + len(versionName)) % 32
    for i in range(exists, 32):
        out.write(b'\x00')
    out.flush()


class ImDictLoader(object):
    def __init__(self) -> None:
        super().__init__()
        self.base = -1
        self.version = 0
        self.root = None
        self.count = 0
        self.path = deque([])

    def load(self, dic_file):
        with open(dic_file, 'rb') as f:
            self._parseHeader(f)
            self._loadRoot(f)
            self._loadOthers(f)
        # for debug
        # with open(dic_file + '.bak', 'wb') as bak:
        #     _writeDictHead(bak)
        #     # base
        #     bak.write(self.base.to_bytes(4, byteorder=ImDict.kByteOrder))
        #     self.root.dump(bak)
        d = ImDict(dic_file, self.base)
        d.useRoot(self.root)
        return d

    def _parseHeader(self, f):
        head = f.read(16).decode()
        if not head.startswith(ImDict.kHead):
            raise IOError('Invalid data file!')
        version = int.from_bytes(f.read(4), byteorder=ImDict.kByteOrder)
        if version > ImDict.kVersionCode:
            raise Exception('Not supported version!')
        self.version = version
        # 28 = 32 - 4
        metadata = f.read(28).decode(encoding=ImDict.kEncoding).strip()

    def _loadRoot(self, f):
        self.base = int.from_bytes(f.read(4), byteorder=ImDict.kByteOrder)
        data = f.read(4)
        if data.decode() == 'root':
            self.root = _Node('', 0, 0)
        else:
            raise Exception('root node is missing!')
        self.count = int.from_bytes(f.read(4), byteorder=ImDict.kByteOrder)
        if self.count <= 0:
            raise Exception('node count error!')

    def _loadOthers(self, f):
        self.path.append(self.root)
        current = self.root
        prevLevel = current.level
        for i in range(0, self.count):
            # 1 node save in 7 bytes
            key = f.read(1).decode()
            offset = int.from_bytes(f.read(4), byteorder=ImDict.kByteOrder)
            level = int.from_bytes(f.read(2), byteorder=ImDict.kByteOrder)
            node = _Node(key, offset, level)
            while len(self.path) > 0 and prevLevel >= level:
                self.path.pop()
                current = self.path[-1]
                prevLevel = current.level
            current.append(node)
            self.path.append(node)
            prevLevel = level
            current = node
